use std::{cell::Cell, ptr::{addr_of, addr_of_mut, null_mut}, sync::Once};
use ctor::ctor;
use libc_print::libc_println;

use super::{mi_heap_is_initialized, mi_thread_id, MiHeap, MiPage, MiPageQueue, MiTld, _mi_heap_collect_abandon, _mi_os_alloc, _mi_os_free};

const MI_LARGE_WSIZE_MAX: usize = 131072;

/// Empty page used to initialize the small free pages array
static _MI_PAGE_EMPTY: MiPage = MiPage::default();

const fn qnull(sz: usize) -> MiPageQueue {
    MiPageQueue {
        first: null_mut(),
        last: null_mut(),
        block_size: sz * std::mem::size_of::<usize>(),
    }
}

const MI_PAGE_QUEUES_EMPTY: [MiPageQueue; 66] = [
    qnull(1),
    qnull(1), qnull(2), qnull(3), qnull(4), qnull(5), qnull(6), qnull(7), qnull(8),
    qnull(10), qnull(12), qnull(14), qnull(16), qnull(20), qnull(24), qnull(28), qnull(32),
    qnull(40), qnull(48), qnull(56), qnull(64), qnull(80), qnull(96), qnull(112), qnull(128),
    qnull(160), qnull(192), qnull(224), qnull(256), qnull(320), qnull(384), qnull(448), qnull(512),
    qnull(640), qnull(768), qnull(896), qnull(1024), qnull(1280), qnull(1536), qnull(1792), qnull(2048),
    qnull(2560), qnull(3072), qnull(3584), qnull(4096), qnull(5120), qnull(6144), qnull(7168), qnull(8192),
    qnull(10240), qnull(12288), qnull(14336), qnull(16384), qnull(20480), qnull(24576), qnull(28672), qnull(32768),
    qnull(40960), qnull(49152), qnull(57344), qnull(65536), qnull(81920), qnull(98304), qnull(114688),
    qnull(MI_LARGE_WSIZE_MAX + 1), // Huge queue
    qnull(MI_LARGE_WSIZE_MAX + 2), // Full queue
];

pub static _MI_HEAP_EMPTY: MiHeap = MiHeap {
    tld: null_mut(),
    pages_free_direct: [addr_of!(_MI_PAGE_EMPTY) as *mut MiPage; 130],
    pages: MI_PAGE_QUEUES_EMPTY,
    thread_delayed_free: std::sync::atomic::AtomicPtr::new(null_mut()),
    thread_id: 0,
    cookie: 0,
    random: 0,
    page_count: 0,
    no_reclaim: false,
};

static _MI_TLD_EMPTY: MiTld = MiTld::default();

struct CleanupGuard<F: FnOnce() + 'static> {
    cleanup: Option<F>
}

impl<F: FnOnce() + 'static> Drop for CleanupGuard<F> {
    fn drop(&mut self) {
        if let Some(cleanup) = self.cleanup.take() {
            cleanup();
        }
    }
}

thread_local! {
    pub static _MI_HEAP_DEFAULT: Cell<*const MiHeap> = Cell::new(&_MI_HEAP_EMPTY as *const MiHeap); 
    static THREAD_CLEANUP: CleanupGuard<fn()> = const {CleanupGuard {
        cleanup: Some(|| {
                mi_thread_done();
        })
    } };
}

pub fn get_heap_default() -> *const MiHeap {
    _MI_HEAP_DEFAULT.with(|h| {h.get()})
}

pub fn set_heap_default(heap: *const MiHeap) {
    _MI_HEAP_DEFAULT.with(|h| {h.set(heap)});
}

static mut _MI_HEAP_MAIN: MiHeap = MiHeap {
    tld: std::ptr::null_mut(),
    pages_free_direct: [addr_of!(_MI_PAGE_EMPTY) as *mut MiPage; 130],
    pages: MI_PAGE_QUEUES_EMPTY,
    thread_delayed_free: std::sync::atomic::AtomicPtr::new(null_mut()),
    thread_id: 0,
    cookie: 0,
    random: 0,
    page_count: 0,
    no_reclaim: false,
};

static mut TLD_MAIN: MiTld = MiTld::default();

static mut _MI_PROCESS_IS_INIT: bool = false;


/* -----------------------------------------------------------
  Initialization and freeing of the thread local heaps
----------------------------------------------------------- */
pub struct MiThreadData {
    pub heap: MiHeap,
    pub tld: MiTld,
}

impl Default for MiThreadData {
    fn default() -> Self {
        Self {
            heap: MiHeap::default(),
            tld: MiTld::default(),
        }
    }
}

fn _mi_is_main_thread() -> bool {
    unsafe {
        _MI_HEAP_MAIN.thread_id == 0 || mi_thread_id() == _MI_HEAP_MAIN.thread_id
    }
}

/// Initialize the thread local default heap, called from `mi_thread_init`
unsafe fn _mi_heap_init() -> bool {
    if mi_heap_is_initialized(get_heap_default()) {return true;}
    if _mi_is_main_thread() {
        // the main heap is statically allocated
        set_heap_default(addr_of!(_MI_HEAP_MAIN));
    } else {
        // use '_mi_os_alloc' to allocate directly from the OS
        let td = _mi_os_alloc(std::mem::size_of::<MiThreadData>()) as *mut MiThreadData;
        if td.is_null() {
            libc_println!("failed to allocate thread local heap memory");
            return false;
        }
        let tld = &mut (*td).tld;
        let heap = &mut (*td).heap;
        std::ptr::copy_nonoverlapping(&_MI_HEAP_EMPTY, heap, 1);
        heap.thread_id = mi_thread_id();
        heap.random = 0x9abcdef0;
        heap.cookie = 0x12345678;
        heap.tld = tld;
        std::ptr::write_bytes(tld, 0, 1);
        tld.heap_backing = heap;
        set_heap_default(heap);
    }
    false
}


/// Free the thread local default heap, called from `mi_thread_done`
unsafe fn _mi_heap_done() -> bool {
    let mut heap = get_heap_default();
    if !mi_heap_is_initialized(heap) {return true;}

    // reset default heap
    set_heap_default(
        if _mi_is_main_thread() {
            addr_of!(_MI_HEAP_MAIN)
        } else {
            &_MI_HEAP_EMPTY
        }
    );

    // switch to backing heap and free it
    heap = (*(*heap).tld).heap_backing;
    if !mi_heap_is_initialized(heap) {return false;}

    // free if not the main thread (or in debug mode)
    if heap != addr_of!(_MI_HEAP_MAIN) {
        if (*heap).page_count > 0 {
            _mi_heap_collect_abandon(heap);
        }
        _mi_os_free(heap as *mut u8, std::mem::size_of::<MiThreadData>());
    }
    false
}


// --------------------------------------------------------
// Try to run `mi_thread_done()` automatically so any memory
// owned by the thread but not yet released can be abandoned
// and re-owned by another thread.
//
// 1. windows dynamic library:
//     call from DllMain on DLL_THREAD_DETACH
// 2. windows static library:
//     use `FlsAlloc` to call a destructor when the thread is done
// 3. unix, pthreads:
//     use a pthread key to call a destructor when a pthread is done
//
// In the last two cases we also need to call `mi_process_init`
// to set up the thread local keys.
// --------------------------------------------------------

/// This is called from the 'mi_malloc_generic`
pub fn mi_thread_init() {
    // ensure our process has started already 
    mi_process_init();

    // initialize the thread local default heap
    if unsafe { _mi_heap_init() } {return;}

    // don't further initialize for the main thread
    if _mi_is_main_thread() {return;}

    THREAD_CLEANUP.with(|_| {});

    libc_println!("thread init: 0x{:x}", mi_thread_id());
}

pub fn mi_thread_done() {
    // abandon the thread local heap
    // return if already ran
    if unsafe { _mi_heap_done() } {return;}

    if !_mi_is_main_thread() {
        libc_println!("thread done: 0x{:x}", mi_thread_id());
    }

}

// --------------------------------------------------------
// Run functions on process init/done, and thread init/done
// --------------------------------------------------------

extern "C" {
    pub fn atexit(f: extern "C" fn()) -> i32;
}

static PROCESS_INIT: Once = Once::new();

fn mi_process_init() {
    unsafe {
        // only initialize once
        PROCESS_INIT.call_once(|| {
            _MI_PROCESS_IS_INIT = true;

            _MI_HEAP_MAIN.tld = addr_of_mut!(TLD_MAIN);
            TLD_MAIN.heap_backing = addr_of_mut!(_MI_HEAP_MAIN);
    
            _MI_HEAP_MAIN.thread_id = mi_thread_id();
            libc_println!("process init: 0x{:x}", _MI_HEAP_MAIN.thread_id);
            _MI_HEAP_MAIN.cookie = 0x12345678;
            _MI_HEAP_MAIN.random = 0x9abcdef0;
            atexit(mi_process_done);
        });
    }
}

static PROCESS_DONE: Once = Once::new();

extern "C" fn mi_process_done() {
    libc_println!("process is shutting down");
    unsafe {
        // only shutdown if we were initialized
        if !_MI_PROCESS_IS_INIT {return;}
        // ensure we are called once
        PROCESS_DONE.call_once(|| {
            libc_println!("process done: 0x{:x}", _MI_HEAP_MAIN.thread_id);
        });

    }
}

#[ctor]
unsafe fn _mi_process_init() {
    mi_process_init();
}


